# db/postgres.py
import psycopg
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.store.postgres.aio import AsyncPostgresStore


class PostgresManager:
    def __init__(
        self,
        dsn: str,
        min_size: int = 1,
        max_size: int = 20,
        timeout: int = 30,
    ):
        # postgresql://postgres:123456@192.168.31.193:5433/llm
        self.dsn = dsn
        self.min_size = min_size
        self.max_size = max_size
        self.timeout = timeout

        self.pool: AsyncConnectionPool | None = None
        self.saver: AsyncPostgresSaver | None = None
        self.store: AsyncPostgresStore | None = None

    async def setup(self):
        """初始化连接池 + LangGraph saver/store"""
        if self.pool:
            return  # 避免重复初始化

        # 使用推荐的方式创建连接池（不在构造函数中打开）
        self.pool = AsyncConnectionPool(
            conninfo=self.dsn,
            min_size=self.min_size,
            max_size=self.max_size,
            timeout=self.timeout,
            kwargs={
                "autocommit": True,  # 避免事务悬挂
                "keepalives_idle": 30,     # TCP keepalive 空闲时间
                "keepalives_interval": 10, # TCP keepalive 探测间隔
                "keepalives_count": 5,     # TCP keepalive 探测次数
            },
            max_idle=300,      # 连接最大空闲时间
            max_lifetime=3600, # 连接最大生存时间
            open=False,        # 不在构造函数中自动打开
        )
        
        # 手动打开连接池
        await self.pool.open()

        # 初始化 LangGraph Saver/Store
        self.saver = AsyncPostgresSaver(self.pool)
        self.store = AsyncPostgresStore(self.pool)

        # 确保表存在
        await self.saver.setup()
        await self.store.setup()

        # 启动时做一次健康检查
        await self.check_health()
        print("✅ PostgresManager 已初始化")

    async def get_saver(self) -> AsyncPostgresSaver:
        """获取PostgreSQL Saver实例"""
        if self.saver is None:
            raise RuntimeError("Postgres saver 未初始化，请先调用setup()")
        return self.saver
    
    async def get_store(self) -> AsyncPostgresStore:
        """获取PostgreSQL Store实例"""
        if self.store is None:
            raise RuntimeError("Postgres store 未初始化，请先调用setup()")
        return self.store

    async def get_connection(self):
        """获取数据库连接（用于自定义查询）"""
        if not self.pool:
            raise RuntimeError("连接池未初始化，请先调用setup()")
        return self.pool.connection()

    async def check_health(self):
        """检查连接池是否可用"""
        if not self.pool:
            raise RuntimeError("连接池未初始化")
        try:
            async with self.pool.connection() as conn:
                await conn.execute("SELECT 1")
        except Exception as e:
            print(f"⚠️ Postgres 健康检查失败: {e}")
            raise

    async def close(self):
        """关闭连接池"""
        if self.pool:
            await self.pool.close()
            self.pool = None
            self.saver = None
            self.store = None
            print("🛑 PostgresManager 已关闭")
